// Copyright (c) 2019 Presto Labs Pte. Ltd.
// Author: jaewon

#include "coin1/fastfeed/double_shift_feed_builder.h"

#include <string>

#include <glog/logging.h>

namespace fastfeed {
namespace detail {

SwitchingContext::SwitchingContext(int num_workers) : last_timestamps_(num_workers, 0) {}

void SwitchingContext::SwitchWorker(int worker_idx, ProtoFeedBuilder* builder) {
  DCHECK_GE(worker_idx, 0);
  DCHECK_LT(worker_idx, last_timestamps_.size());

  if (active_builder_) {
    active_builder_->DisableUserCallback();
  }
  active_worker_idx_ = worker_idx;
  active_builder_ = builder;
  active_builder_->EnableUserCallback();
}

bool SwitchingContext::UpdateAndCheck(int worker_idx, int64_t timestamp) {
  DCHECK_GE(worker_idx, 0);
  DCHECK_LT(worker_idx, last_timestamps_.size());

  last_timestamps_[worker_idx] = timestamp;
  if (active_worker_idx_ == -1) {
    return true;
  }
  if (active_worker_idx_ == worker_idx) {
    return false;
  }
  auto diff = timestamp - last_timestamps_[active_worker_idx_];
  return diff >= threshold_;
}

}  // namespace detail

DoubleShiftFeedBuilder::DoubleShiftFeedBuilder(int num_workers) { InitBuilders(num_workers); }

DoubleShiftFeedBuilder::DoubleShiftFeedBuilder(int num_workers, const std::string& sub_req)
    : FeedBuilderBase(sub_req) {
  InitBuilders(num_workers);
}

void DoubleShiftFeedBuilder::InitBuilders(int num_workers) {
  CHECK(builders_.empty());
  for (int i = 0; i < num_workers; ++i) {
    auto builder = new ProtoFeedBuilder(sub_req());
    builder->DisableUserCallback();
    builders_.emplace_back(builder);
  }
}

void DoubleShiftFeedBuilder::set_from_timestamp(std::optional<int64_t> from_timestamp) {
  for (auto& builder : builders_) {
    builder->set_from_timestamp(from_timestamp);
  }
}

std::function<void(const proto::Feed&)> DoubleShiftFeedBuilder::GetUpdateFunction(int worker_idx) {
  return std::bind(
      &DoubleShiftFeedBuilder::UpdateByProto,
      this,
      worker_idx,
      builders_[worker_idx].get(),
      std::placeholders::_1);
}

void DoubleShiftFeedBuilder::SetSymbolGroup(
    const std::string& symbol,
    const std::string& symbol_group) {
  auto it = symbol_to_group_.find(symbol);
  if (it == symbol_to_group_.end()) {
    symbol_to_group_[symbol] = symbol_group;
    return;
  }
  CHECK_EQ(it->second, symbol_group);
  return;
}

const std::string& DoubleShiftFeedBuilder::FindSymbolGroup(const std::string& symbol) {
  auto it = symbol_to_group_.find(symbol);
  if (it != symbol_to_group_.end()) {
    return it->second;
  }
  symbol_to_group_[symbol] = symbol;
  return symbol;
}

detail::SwitchingContext* DoubleShiftFeedBuilder::FindSwitchingContextBySymbol(
    const std::string& symbol) {
  {
    auto it = symbol_to_switching_ctx_.find(symbol);
    if (it != symbol_to_switching_ctx_.end()) {
      return it->second;
    }
  }

  {
    const std::string& symbol_group = FindSymbolGroup(symbol);
    detail::SwitchingContext* ctx = nullptr;

    auto it = switching_ctxs_.find(symbol_group);
    if (it != switching_ctxs_.end()) {
      ctx = it->second.get();
    } else {
      ctx = new detail::SwitchingContext(builders_.size());
      switching_ctxs_[symbol_group].reset(ctx);
    }
    symbol_to_switching_ctx_[symbol] = ctx;
    return ctx;
  }
}

void DoubleShiftFeedBuilder::UpdateByProto(
    int worker_idx,
    ProtoFeedBuilder* builder,
    const proto::Feed& feed) {
  if (feed.feed_type() == proto::FEED_TYPE_METADATA &&
      feed.data_type() == proto::DATA_TYPE_FEED_ARB_PARAMETER) {
    CHECK(feed.has_feed_arb_param());
    SetSymbolGroup(feed.symbol(), feed.feed_arb_param().symbol_group());
    return;
  }

  // Overhead: aymptotically adding one std::unordered_map search.
  auto switching_ctx = FindSwitchingContextBySymbol(feed.symbol());
  if (switching_ctx->UpdateAndCheck(worker_idx, feed.timestamp())) {
    switching_ctx->SwitchWorker(worker_idx, builder);
  }
  builder->UpdateByProto(feed);
}

void DoubleShiftFeedBuilder::Subscribe(
    proto::FeedType feed_type,
    const std::string& custom_feed_type,
    const std::string& symbol,
    const std::string& view_type,
    std::function<void(const view::ViewBase*)> callback) {
  for (auto& builder : builders_) {
    builder->Subscribe(feed_type, custom_feed_type, symbol, view_type, callback);
  }
}

}  // namespace fastfeed
